
#include "lp_norm_v2_tiling.h"
#include "math.h"
#include "register/op_def_registry.h"
#include "tiling_common.h"
#include <cstdint>

namespace optiling
{
static ge::graphStatus TilingFunc(gert::TilingContext *context)
{
    LpNormV2TilingData tiling;

    // 设置size
    auto x_shape = context->GetInputShape(0);
    auto x_orgin_shape = x_shape->GetStorageShape();
    auto data_size = x_orgin_shape.GetShapeSize();
    auto dtype = context->GetInputDesc(0)->GetDataType();
    auto dtypeSize = GetSizeByDataType(dtype);
    tiling.set_size(data_size);
    // 设置参数
    auto runtime_attrs = context->GetAttrs();
    auto p_value = *(runtime_attrs->GetFloat(0));
    tiling.set_p(p_value);
    tiling.set_epsilon(*(runtime_attrs->GetFloat(3)));

    // 定义 sort fields
    int64_t sortByDimPreSize = 0;
    int64_t sortByDimDimSize = 0;
    int64_t sortByDimDimBlockLength = 0;
    int64_t sortByDimDimLength = 0;
    int64_t workspace_size = 0;
    // 设置 目标 dimSize 和 dim length
    // 获取 axes选项
    auto axes = runtime_attrs->GetListInt(1);
    // 配置 sort fields
    while (axes->GetSize() != 0)
    {
        auto target_dim_idx = axes->GetData()[0];
        std::cout << "xxxxxxxxxxx" << target_dim_idx << std::endl;
        auto dim_num = x_orgin_shape.GetDimNum();
        // 判断参数是否正确
        if (target_dim_idx < 0 || target_dim_idx >= dim_num)
        {
            break;
        }
        // 计算前置数量
        sortByDimPreSize = 1;
        for (auto i = 0; i < target_dim_idx; ++i)
        {
            sortByDimPreSize *= x_orgin_shape.GetDim(i);
        }
        // 获取目标维度大小
        sortByDimDimSize = x_orgin_shape.GetDim(target_dim_idx);
        // 这里是与最后一个维度做sort, 所以blocklength = 1 * dataSize。即一个数据为一个单位。
        sortByDimDimBlockLength = dtypeSize;
        // 获取目标目标维度之间的维度乘积
        sortByDimDimLength = 1;
        for (auto i = target_dim_idx + 1; i < dim_num; ++i)
        {
            sortByDimDimLength *= x_orgin_shape.GetDim(i);
        }
        break;
    }
    tiling.set_sortByDimPreSize(sortByDimPreSize);
    tiling.set_sortByDimDimSize(sortByDimDimSize);
    tiling.set_sortByDimDimLength(sortByDimDimLength);
    tiling.set_sortByDimDimBlockLength(sortByDimDimBlockLength);

    // 如果需要transpose, 分割的方法与整体运算不同
    auto vector_align_size = 32;
    auto mem_addition_num = 3;
    if (p_value == 0)
    {
        // 因为需要使用Compare所以 使用256对齐
        vector_align_size = 256;
        mem_addition_num = 5;
    }
    else if (p_value == 1)
    {
    }
    else if (p_value == float(1.0) / float(0.0))
    {
    }
    else if (p_value == float(-1.0) / float(0.0))
    {
    }
    else
    {
        switch (dtype)
        {
        case ge::DataType::DT_FLOAT16: {
            mem_addition_num = 8;
        }
        break;
        default:
            break;
        }
    }
    if (sortByDimDimSize != 0)
    {
        auto blockDim = common_tiling_set_kernel_vec(tiling, context);
        auto alignCount = data_size / sortByDimDimSize;
        // 分配运算个数
        COMMON_TILING_SPLIT((&tiling), blockDim, alignCount, 1);

        // 对一次求解的数据进行分割，算出每次运算多少
        auto alignDataSize = CEIL_DIV(vector_align_size, dtypeSize);
        auto alignDimSize = ALIGN_TO(sortByDimDimSize, alignDataSize);
        auto total_count = (1 + 1) * 2 + mem_addition_num;
        auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
        auto tileLength =
            common_tiling_auto_tileLenth(ascendcPlatform, dtypeSize, alignDataSize, total_count, alignDimSize);
        tiling.set_formerTileLength(tileLength);
        tiling.set_tailTileLength(tileLength);
        // 计算最后一次循环
        int64_t finalCalcCount = ALIGN_TO((alignDimSize % tileLength), alignDataSize);
        if (finalCalcCount == 0)
        {
            finalCalcCount = tileLength;
        }
        tiling.set_formerFinalCalcCount(finalCalcCount);
        tiling.set_tailFinalCalcCount(finalCalcCount);
        tiling.set_dType(dtype);
    }
    else
    {
        CommonTiling<LpNormV2TilingData> commonTiling;
        commonTiling.tiling = &tiling;
        commonTiling.context = context;
        commonTiling.data_size = data_size;
        commonTiling.mem_addition_num = mem_addition_num;
        commonTiling.vector_align_size = vector_align_size;
        commonTiling.apply();
    }
    // 序列化
    tiling_save_to_buffer(tiling, context);
    tiling_zero_workspace(context);
    print_common_tiling(tiling);
    std::cout << "axes->GetSize()" << axes->GetSize() << std::endl;
    std::cout << "sortByDimPreSize" << sortByDimPreSize << std::endl;
    std::cout << "sortByDimDimSize" << sortByDimDimSize << std::endl;
    std::cout << "sortByDimDimLength" << sortByDimDimLength << std::endl;
    std::cout << "sortByDimDimBlockLength" << sortByDimDimBlockLength << std::endl;
    return ge::GRAPH_SUCCESS;
}
} // namespace optiling

namespace ge
{
static ge::graphStatus InferShape(gert::InferShapeContext *context)
{
    std::cout << "InferShape" << std::endl;
    gert::Shape *y_shape = context->GetOutputShape(0);
    y_shape->SetDim(0, 1);
    y_shape->SetDimNum(1);

    auto runtime_attrs = context->GetAttrs();
    auto axes = runtime_attrs->GetListInt(1);
    // 如果选定某一个维度，则输出为 size / 该维度
    while (axes->GetSize() != 0)
    {
        const gert::Shape *x1_shape = context->GetInputShape(0);
        auto target_dim_idx = axes->GetData()[0];
        auto dim_num = x1_shape->GetDimNum();
        // 判断参数是否正确
        if (target_dim_idx < 0 || target_dim_idx >= dim_num)
        {
            break;
        }

        auto dim_size = x1_shape->GetDim(target_dim_idx);
        y_shape->SetDim(0, x1_shape->GetShapeSize() / dim_size);
        break;
    }

    return GRAPH_SUCCESS;
}
} // namespace ge

namespace ops
{
class LpNormV2 : public OpDef
{
  public:
    explicit LpNormV2(const char *name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("p").AttrType(OPTIONAL).Float(2);
        this->Attr("axes").AttrType(OPTIONAL).ListInt({});
        this->Attr("keepdim").AttrType(OPTIONAL).Bool(false);
        this->Attr("epsilon").AttrType(OPTIONAL).Float(1e-12);

        this->SetInferShape(ge::InferShape);

        this->AICore().SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");
    }
};

OP_ADD(LpNormV2);
} // namespace ops
